#!/usr/bin/python
# coding = utf-8
# by cosine
# cos@knownsec.com

"""
线程池初始化：

import threadpool

poolSize = 10 # 线程池大小
poolModel = 0 # 线程池模式
tp = threadpool.ThreadPool(poolSize,poolModel)

具体示例看该程序if __name__ == '__main__':部分


=== ChangeLog ===
2011/01/14
修补<模式2>几处小bugs：
增加getDeadWorkersNumber函数替代之前用threading.activeCount()来统计不活跃线程worker实例
这样更加精确，感谢yc同学反馈
增加线程池初始化的参数容错判断，感谢watercloud反馈

2010/12/31
线程池模式增加<模式4>
类似<模式3>，唯一差异性看waitforComplete函数里相关代码注释
仔细整理一遍该模块：
1、增加Message消息机制
2、整理相关注释、实例等

2010/09/15
线程池模式增加<模式3>
类似<模式1>：队列的实际大小随时间而变化，每个线程所处理的队列任务在任务结束时会发送task_done信号，如果队列没任务时，线程池退出
但此时退出消息机制来自外部通知，这样会更加灵活

2010/02/10
当线程池模式为0时，修补子线程不完美退出bugs

2009/09/03
线程池支持三种模式：
<模式0>
队列的实际大小初始化后不会再改变
每个线程最后都会从队列里接收到线程退出消息{'exit_flag':'exit threadpool'}，线程池退出
此时退出消息机制来自内部，不需要外部发送退出消息
<模式1>
队列的实际大小随时间而变化
每个线程所处理的队列任务在任务结束时会发送task_done信号，如果队列没任务时，线程池退出
此时退出消息机制来自内部，不需要外部发送退出消息
<模式2>
线程池将守护着永远不退出

2009/07/06
线程池第一个稳定版本从石头蹦出


"""

import Queue
import threading
import sys

class Message:
	'''
	消息通知类，由ThreadPool实例化
	ThreadPool通过该消息实例来管理每个Worker的退出等消息通知
	'''
	pass

class Worker(threading.Thread):
	'''线程类'''
	def __init__(self, poolModel, taskQueue, message):
		'''init a new thread'''
		threading.Thread.__init__(self)
		self.poolModel = poolModel
		self.taskQueue = taskQueue
		self.message = message
		self.start()

	def run(self):
		'''
		线程运行：
		callback是一个函数名
		argv是这个函数的参数字典
		'''
		while True:
			try:
				if self.poolModel == 0:
					callback, argv = self.taskQueue.get()
					if argv.has_key('exit_flag'): # <模式0>的退出消息
						if argv['exit_flag'] == 'exit threadpool':
							self.taskQueue.task_done() # 任务发送task_done信号，表示任务运行结束
							break
						
				elif self.poolModel == 1:
					if self.message.threadexit: # <模式1>的退出消息
						break
					try:
						callback, argv = self.taskQueue.get(True,self.message.queue_timeout)
					except Queue.Empty:
						continue
					
				elif self.poolModel == 2:
					# <模式2>永远不退出，非常霸道，具体看ThreadPool类里的waitforComplete函数<模式2>相关代码
					callback, argv = self.taskQueue.get()
					
				elif self.poolModel in (3,4): # <模式3><模式4>
					try:
						callback, argv = self.taskQueue.get(True,self.message.queue_timeout)
					except Queue.Empty:
						if not self.taskQueue.unfinished_tasks and self.message.threadexit:
							# <模式3><模式4>的退出消息
							break
						continue
				

				callback(argv) # 执行回调函数
				self.taskQueue.task_done() # 任务发送task_done信号，表示任务运行结束
			except Queue.Empty:
				pass

class ThreadPool:
	'''线程池管理类'''
	def __init__(self, poolSize, poolModel=2):
		poolModelList = [0,1,2,3,4]
		if poolModel not in poolModelList:
			sys.exit('value of poolModel must be in %s'%poolModelList)
		try:
			poolSize = int(poolSize)
		except Exception,e:
			sys.exit('value of poolSize is not a valid integer')

		# 初始化消息实例
		self.message = Message()
		self.message.threadexit = False
		self.message.queue_timeout = 3

		self.taskQueue = Queue.Queue(0)
		self.workers = []
		self.poolSize = poolSize
		self.poolModel = poolModel
		self.createWorkers(poolSize, poolModel)
	
	def createWorkers(self, poolSize, poolModel):
		'''创建一个新的worker线程'''
		for i in xrange(poolSize):
			worker = Worker(poolModel, self.taskQueue, self.message)
			self.workers.append(worker)

	def waitforComplete(self):
		'''等待线程池任务结束'''
		if self.poolModel == 0:
			for i in xrange(self.poolSize):
				self.putToQueue(None,exit_flag='exit threadpool')
			while len(self.workers):
				# 修补子线程不完美退出，2010-02-10
				worker = self.workers.pop()
				worker.join()

		if self.poolModel == 1:
			self.taskQueue.join()
			self.message.threadexit = True

		if self.poolModel == 2:
			# 检查工作队列长度，去掉死掉的线程，并保持线程池worker个数饱满。
			while len(self.workers):
				worker = self.workers.pop()
				worker.join()
				#if threading.activeCount() < self.poolSize:
				deadWorkersNum = self.getDeadWorkersNumber()
				if deadWorkersNum:
					for w in self.workers:
						if not w.isAlive():
							self.workers.remove(w) # 移除已死线程worker实例
					self.createWorkers(deadWorkersNum,self.poolModel)

		if self.poolModel == 3:
			# 对线程进行join
			while len(self.workers):
				worker = self.workers.pop()
				worker.join()

		if self.poolModel == 4:
			# 对任务队列进行join
			self.taskQueue.join()

	def putToQueue(self, callback, **argv):
		'''将任务填入线程池队列中，待分配'''
		self.taskQueue.put((callback, argv))

	def getJobsNumber(self):
		'''获取线程池中taskQueue还未分配的队列大小'''
		return self.taskQueue.qsize()

	def getUnfinishedNumber(self):
		'''
		获取线程池未完成任务的任务数
		
		特别注意：
		该线程池中getJobsNumber()与getUnfinishedNumber()的区别：
		1、getJobsNumber()表示线程池中taskQueue还未分配的队列大小
		2、getUnfinishedNumber()表示线程池未完成任务的任务数
		3、getUnfinishedNumber() >= getJobsNumber()
		'''
		return self.taskQueue.unfinished_tasks
	
	def getDeadWorkersNumber(self):
		'''获取不活跃的线程worker实例'''
		deadWorkersNum = 0
		for w in self.workers:
			if not w.isAlive():
				deadWorkersNum += 1
		return deadWorkersNum

	def popAllJobs(self):
		'''弹出taskQueue中的所有待分配的任务，返回任务list'''
		results = []
		qsize = self.taskQueue.qsize()
		i = 0
		while i < qsize*2:
			try:
				item = self.taskQueue.get_nowait()
				results.append(item)
				i += 1
			except :
				return results
		return results
	
	def setExitFlag(self):
		'''设置线程池结束标志'''
		self.message.threadexit = True
		self.message.queue_timeout = 1

if __name__ == '__main__':
	# demo for threadpool
	import time

	def stat():
		'''统计线程池中getJobsNumber()与getUnfinishedNumber()的大小'''
		while 1:
			jobsNumber = tp.getJobsNumber()
			unfinishedNumber = tp.getUnfinishedNumber()
			statVaue = '>> jobs undistributed number: %s, jobs unfinished number: %s'%(jobsNumber,unfinishedNumber)
			print statVaue
			time.sleep(1)
			if tp.getJobsNumber() == 0 and tp.getUnfinishedNumber() == 0 and poolModel in (3,4):
				tp.setExitFlag()
				print '>> send thread exit message...'

	def func(args):
		value1 = args['arg1']
		value2 = args['arg2']
		if value1:
			msg = '[+] value1: %s, value2: %s'%(value1,value2)
			print msg
			time.sleep(2)
			tp.putToQueue(func,arg1=value1-1,arg2=value2)
		else:
			msg = '[-] value1: %s, value2: %s'%(value1,value2)
			print msg

	poolSize = 10
	poolModel = 3 # 可以修改该模式值查看不同模式的效果
	tp = ThreadPool(poolSize,poolModel)

	# 单独线程开启stat统计
	t=threading.Thread(target=stat, args=())
	t.setDaemon(True)
	t.start()

	value1=1
	for i in xrange(20):
		value2 = i
		tp.putToQueue(func,arg1=value1,arg2=value2)
	tp.waitforComplete()
	print 'threadpool is dry:D'


